import torch
from torch.nn import Module, Conv2d, MaxPool2d


class SharedInception(Module):
    def __init__(self, input_dim):
        super(SharedInception, self).__init__()
        self.conv1 = Conv2d(input_dim, input_dim, kernel_size=(1, 1))
        self.conv3 = Conv2d(input_dim, input_dim, kernel_size=(3, 3), padding=1)
        self.conv5 = Conv2d(input_dim, input_dim, kernel_size=(5, 5), padding=2)
        self.pool = MaxPool2d(kernel_size=(3, 3), stride=1, padding=1)

    def forward(self, x):
        res_list = [self.conv1(x), self.conv3(x), self.conv5(x), self.pool(x)]
        res = torch.stack(res_list, dim=-1).mean(-1)
        return res
